﻿/*
BF4Rcon.NET, provides a .NET implementation of BF4's RCON interface.
Copyright (C) 2013 Tim Marro

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace TimScripts.BF4Rcon.Frostbite
{
    public class RequestManager : IRequestManager, IDisposable
    {
        private Dictionary<int, Packet> Responses = new Dictionary<int, Packet>();
        private Dictionary<int, TaskCompletionSource<Packet>> ResponseTasks = new Dictionary<int, TaskCompletionSource<Packet>>();
        private Dictionary<int, ManualResetEvent> ManualResetEvents = new Dictionary<int, ManualResetEvent>();
        bool IsClosed = false;
        Exception closeException = null;
        private object ResponsesLock = new object();
        private bool Disposed;
        //private object MreLock = new object();

        public Task<Packet> GetResponseAsync(int sequence)
        {
            return GetResponseAsync(sequence, CancellationToken.None);
        }

        public Task<Packet> GetResponseAsync(int sequence, CancellationToken ct)
        {
            Packet response;
            TaskCompletionSource<Packet> tcs;

            lock (ResponsesLock)
            {
                if (IsClosed)
                    throw closeException;

                if (Responses.TryGetValue(sequence, out response))
                    return TaskEx.FromResult(response);

                tcs = new TaskCompletionSource<Packet>();
                ResponseTasks.Add(sequence, tcs);

                ct.Register(() =>
                {
                    tcs.TrySetCanceled();
                });

                return tcs.Task;
            }
        }

        public Packet GetResponse(int sequence)
        {
            return GetResponse(sequence, Timeout.Infinite);
        }

        public Packet GetResponse(int sequence, int timeout)
        {
            Packet response;
            ManualResetEvent mre;

            lock (ResponsesLock)
            {
                if (IsClosed)
                    throw closeException;

                if (Responses.TryGetValue(sequence, out response))
                {
                    Responses.Remove(sequence);
                    return response;
                }

                mre = new ManualResetEvent(false);
                ManualResetEvents[sequence] = mre;
            }

            mre.WaitOne(timeout);
            mre.Dispose();

            lock (ResponsesLock)
            {
                ManualResetEvents.Remove(sequence);

                if (Responses.TryGetValue(sequence, out response))
                {
                    Responses.Remove(sequence);
                    return response;
                }

                if (IsClosed)
                    throw closeException;

                //timeout at this point
                throw new TimeoutException();
            }
        }

        public void AddResponse(Packet packet)
        {
            if (packet == null)
                throw new ArgumentNullException("packet");

            int seq = packet.Sequence;
            TaskCompletionSource<Packet> tcs;
            ManualResetEvent mre;

            lock (ResponsesLock)
            {
                if (ResponseTasks.TryGetValue(seq, out tcs))
                {
                    Responses.Remove(seq);
                    ResponseTasks.Remove(seq);
                    tcs.TrySetResult(packet);
                }
                else
                {
                    Responses[seq] = packet;

                    if (ManualResetEvents.TryGetValue(seq, out mre))
                    {
                        mre.Set();
                    }
                }
            }
        }

        public void Close(Exception exception)
        {
            lock (ResponsesLock)
            {
                closeException = exception;
                IsClosed = true;

                foreach (TaskCompletionSource<Packet> tcs in ResponseTasks.Values)
                {
                    tcs.TrySetException(exception);
                }

                foreach (ManualResetEvent mre in ManualResetEvents.Values)
                {
                    mre.Set();
                }

                ResponseTasks.Clear();
                ManualResetEvents.Clear();
            }
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!Disposed)
            {
                if (disposing)
                {
                    Close(new ObjectDisposedException(null));

                    Disposed = true;
                }
            }
        }
    }
}
